"""
Author: Fei Wu
Adapted by Danni Wu
Date: 0714
"""
import os
import numpy as np

import cv2
from PIL import Image
from sklearn.decomposition import PCA
import mmcv
import torch

def visual_feature(feature,layer="img",path="runs/wf_feat_analysis/lidar_only"):
    print("feature shape: ",feature.shape)

    def reshape_normalize(x):
        """
        Args: 
            x : [B,C,H,W]
        """
        B, C, H, W = x.shape
        x = x.transpose([0, 2, 3, 1]).reshape([-1, C])

        denominator = np.linalg.norm(x, axis=-1, keepdims=True)
        denominator = np.where(denominator==0, 1, denominator)
        return x / denominator

    def normalize(x):
        denominator = np.linalg.norm(x, axis=-1, keepdims=True)
        denominator = np.where(denominator == 0, 1, denominator)
        return x / denominator

    grd_feat = feature.data.cpu().numpy() 
    B, C, H, W = grd_feat.shape
    
    pca = PCA(n_components=3)
    pca.fit(reshape_normalize(grd_feat))
    
    grd_feat=((normalize(pca.transform(reshape_normalize(grd_feat))) + 1) / 2).reshape(B, H, W, 3)
    # print("grd_feat: ",grd_feat.shape,grd_feat[0].shape)

    mmcv.mkdir_or_exist(path)
        
    for b in range(B):
        grd = Image.fromarray((grd_feat[b] * 255).astype(np.uint8))
        grd = grd.resize((W, H))
        print(f"{layer} saving to.... {path}/{layer}_{b}.png")
        grd.save(f'{path}/{layer}_{b}.png')
        
        org = cv2.imread(f'{path}/{layer}_{b}.png')
        org = cv2.flip(org,-1)
        cv2.imwrite(f'{path}/{layer}_{b}.png',org)
        
    # # grd.save(path+layer+".png")
    # count=0
    # while os.path.exists(f"{path}/{layer}.png"):
    #     layer=layer+str(count)
    #     count=count+1

    # print(f"saving feature.............: {path}/{layer}.png")
    # grd.save(f"{path}/{layer}.png")
    
    
    
def save_depth_image(img, tag="depth_img", path="feats_vis"):
    """
    depth image: [B, N, C, H, W]
    """
    if not os.path.exists(path):
        os.mkdir(path)
        
    img = img.permute(0, 1, 3, 4, 2) #[B, N, H, W, C]
    B, N, H, W, C = img.shape
        
    for b in range(B):
        for i in range(N):
            image_depth = img[b][i].cpu().detach().numpy()
            # print(image)
            # print(image.shape)
            image_color = cv2.applyColorMap(cv2.convertScaleAbs(image_depth, alpha=15), cv2.COLORMAP_JET)
            img_vis = Image.fromarray(image_color)
            print(f"{tag} saving to.... {path}/{tag}_{b}_{i}.png")
            img_vis.save(f'{path}/{tag}_{b}_{i}.png')
            
            org = cv2.imread(f'{path}/{tag}_{b}_{i}.png')
            org = cv2.flip(org,-1)
            cv2.imwrite(f'{path}/{tag}_{b}_{i}.png',org)

def denorm_tensor_img(img:torch.Tensor):
    img_np = img.permute(1,2,0).detach().cpu().numpy()
    img_np = (img_np-img_np.min())/(img_np.max()-img_np.min())*255
    return img_np

def plot_boxes_on_image(image, boxes, output_path):
    """
    Plots bounding boxes on an image and saves the result.
    
    Parameters:
    - image_path (str): Path to the input image.
    - boxes (numpy.ndarray): An array of shape (N, 4) containing the bounding boxes (x1, y1, x2, y2).
    - output_path (str): Path to save the output image.
    """
    import matplotlib.pyplot as plt
    # Load the image

    # Create a figure and axis
    fig, ax = plt.subplots(1)
    
    # Display the image
    ax.imshow(image)

    # Plot each box
    for box in boxes:
        x1, y1, x2, y2 = box
        width = x2 - x1
        height = y2 - y1
        rect = plt.Rectangle((x1, y1), width, height, edgecolor='red', facecolor='none', linewidth=2)
        ax.add_patch(rect)

    # Hide the axes
    ax.axis('off')

    # Save the result
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
    plt.close(fig)